"""Code parser using tree-sitter for AST analysis."""

from pathlib import Path
from typing import Dict, List, Optional, Any
import tree_sitter
from tree_sitter import Language, Parser
import tree_sitter_python as tspython
import tree_sitter_javascript as tsjavascript
import tree_sitter_typescript as tstypescript
import tree_sitter_go as tsgo
import tree_sitter_java as tsjava

class CodeParser:
    """Parse source code using tree-sitter for AST analysis."""
    
    def __init__(self):
        self.parsers = {}
        self._initialize_parsers()
    
    def _initialize_parsers(self):
        """Initialize tree-sitter parsers for supported languages."""
        # Build languages
        PYTHON = Language(tspython.language(), "python")
        JAVASCRIPT = Language(tsjavascript.language(), "javascript")
        TYPESCRIPT = Language(tstypescript.language_typescript(), "typescript")
        TSX = Language(tstypescript.language_tsx(), "tsx")
        GO = Language(tsgo.language(), "go")
        JAVA = Language(tsjava.language(), "java")
        
        # Initialize parsers
        self.parsers = {
            '.py': self._create_parser(PYTHON),
            '.js': self._create_parser(JAVASCRIPT),
            '.ts': self._create_parser(TYPESCRIPT),
            '.tsx': self._create_parser(TSX),
            '.jsx': self._create_parser(JAVASCRIPT),
            '.go': self._create_parser(GO),
            '.java': self._create_parser(JAVA),
        }
    
    def _create_parser(self, language) -> Parser:
        """Create a parser for a given language."""
        parser = Parser()
        parser.set_language(language)
        return parser
    
    def parse_file(self, file_path: Path) -> Optional[Dict[str, Any]]:
        """Parse a source file and extract structured information."""
        extension = file_path.suffix
        if extension not in self.parsers:
            return None
        
        try:
            content = file_path.read_text(encoding='utf-8')
            parser = self.parsers[extension]
            tree = parser.parse(bytes(content, 'utf8'))
            
            if extension == '.py':
                return self._parse_python(tree, content)
            elif extension in ['.js', '.ts', '.tsx', '.jsx']:
                return self._parse_javascript(tree, content)
            elif extension == '.go':
                return self._parse_go(tree, content)
            elif extension == '.java':
                return self._parse_java(tree, content)
            
        except Exception as e:
            return None
    
    def _parse_python(self, tree, content: str) -> Dict[str, Any]:
        """Parse Python source code."""
        root_node = tree.root_node
        
        functions = []
        classes = []
        imports = []
        
        def walk_tree(node):
            if node.type == 'function_definition':
                func_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_docstring(node, content)
                params = self._get_python_params(node)
                start_line = node.start_point[0] + 1
                end_line = node.end_point[0] + 1
                
                functions.append({
                    'name': func_name,
                    'type': 'function',
                    'docstring': docstring,
                    'params': params,
                    'start_line': start_line,
                    'end_line': end_line,
                    'language': 'python'
                })
            
            elif node.type == 'class_definition':
                class_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_docstring(node, content)
                methods = []
                
                # Find methods in class
                for child in node.children:
                    if child.type == 'block':
                        for method_node in child.children:
                            if method_node.type == 'function_definition':
                                method_name = self._get_node_text(method_node.child_by_field_name('name'), content)
                                method_docstring = self._get_docstring(method_node, content)
                                methods.append({
                                    'name': method_name,
                                    'docstring': method_docstring
                                })
                
                classes.append({
                    'name': class_name,
                    'type': 'class',
                    'docstring': docstring,
                    'methods': methods,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'python'
                })
            
            elif node.type == 'import_statement' or node.type == 'import_from_statement':
                import_text = self._get_node_text(node, content)
                imports.append({
                    'statement': import_text,
                    'type': 'import',
                    'line': node.start_point[0] + 1
                })
            
            for child in node.children:
                walk_tree(child)
        
        walk_tree(root_node)
        
        return {
            'functions': functions,
            'classes': classes,
            'imports': imports,
            'language': 'python'
        }
    
    def _parse_javascript(self, tree, content: str) -> Dict[str, Any]:
        """Parse JavaScript/TypeScript source code."""
        root_node = tree.root_node
        
        functions = []
        classes = []
        imports = []
        
        def walk_tree(node):
            if node.type == 'function_declaration' or node.type == 'arrow_function':
                func_name = self._get_node_text(node.child_by_field_name('name'), content)
                if not func_name and node.parent and node.parent.type == 'variable_declarator':
                    func_name = self._get_node_text(node.parent.child_by_field_name('name'), content)
                
                docstring = self._get_jsdoc(node, content)
                params = self._get_js_params(node)
                
                functions.append({
                    'name': func_name or 'anonymous',
                    'type': 'function',
                    'docstring': docstring,
                    'params': params,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'javascript'
                })
            
            elif node.type == 'class_declaration':
                class_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_jsdoc(node, content)
                methods = []
                
                # Find methods in class
                for child in node.children:
                    if child.type == 'class_body':
                        for method_node in child.children:
                            if method_node.type == 'method_definition':
                                method_name = self._get_node_text(method_node.child_by_field_name('name'), content)
                                method_docstring = self._get_jsdoc(method_node, content)
                                methods.append({
                                    'name': method_name,
                                    'docstring': method_docstring
                                })
                
                classes.append({
                    'name': class_name,
                    'type': 'class',
                    'docstring': docstring,
                    'methods': methods,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'javascript'
                })
            
            elif node.type == 'import_statement' or node.type == 'import_declaration':
                import_text = self._get_node_text(node, content)
                imports.append({
                    'statement': import_text,
                    'type': 'import',
                    'line': node.start_point[0] + 1
                })
            
            for child in node.children:
                walk_tree(child)
        
        walk_tree(root_node)
        
        return {
            'functions': functions,
            'classes': classes,
            'imports': imports,
            'language': 'javascript'
        }
    
    def _parse_go(self, tree, content: str) -> Dict[str, Any]:
        """Parse Go source code."""
        root_node = tree.root_node
        
        functions = []
        types = []
        imports = []
        
        def walk_tree(node):
            if node.type == 'function_declaration':
                func_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_go_comment(node, content)
                params = self._get_go_params(node)
                
                functions.append({
                    'name': func_name,
                    'type': 'function',
                    'docstring': docstring,
                    'params': params,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'go'
                })
            
            elif node.type == 'type_declaration':
                type_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_go_comment(node, content)
                
                types.append({
                    'name': type_name,
                    'type': 'type',
                    'docstring': docstring,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'go'
                })
            
            elif node.type == 'import_declaration':
                import_text = self._get_node_text(node, content)
                imports.append({
                    'statement': import_text,
                    'type': 'import',
                    'line': node.start_point[0] + 1
                })
            
            for child in node.children:
                walk_tree(child)
        
        walk_tree(root_node)
        
        return {
            'functions': functions,
            'classes': types,
            'imports': imports,
            'language': 'go'
        }
    
    def _parse_java(self, tree, content: str) -> Dict[str, Any]:
        """Parse Java source code."""
        root_node = tree.root_node
        
        methods = []
        classes = []
        imports = []
        
        def walk_tree(node):
            if node.type == 'method_declaration':
                method_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_java_comment(node, content)
                params = self._get_java_params(node)
                
                methods.append({
                    'name': method_name,
                    'type': 'method',
                    'docstring': docstring,
                    'params': params,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'java'
                })
            
            elif node.type == 'class_declaration':
                class_name = self._get_node_text(node.child_by_field_name('name'), content)
                docstring = self._get_java_comment(node, content)
                
                classes.append({
                    'name': class_name,
                    'type': 'class',
                    'docstring': docstring,
                    'start_line': node.start_point[0] + 1,
                    'end_line': node.end_point[0] + 1,
                    'language': 'java'
                })
            
            elif node.type == 'import_declaration':
                import_text = self._get_node_text(node, content)
                imports.append({
                    'statement': import_text,
                    'type': 'import',
                    'line': node.start_point[0] + 1
                })
            
            for child in node.children:
                walk_tree(child)
        
        walk_tree(root_node)
        
        return {
            'functions': methods,
            'classes': classes,
            'imports': imports,
            'language': 'java'
        }
    
    def _get_node_text(self, node, content: str) -> str:
        """Get text content from a node."""
        if not node:
            return ""
        return content[node.start_byte:node.end_byte].decode('utf8')
    
    def _get_docstring(self, function_node, content: str) -> str:
        """Extract docstring from a function."""
        for child in function_node.children:
            if child.type == 'block':
                for stmt in child.children:
                    if stmt.type == 'expression_statement':
                        expr = stmt.children[0]
                        if expr.type == 'string':
                            return self._get_node_text(expr, content).strip('"\'')
        return ""
    
    def _get_python_params(self, function_node) -> List[Dict[str, str]]:
        """Extract Python function parameters."""
        params = []
        params_node = function_node.child_by_field_name('parameters')
        if params_node:
            for param in params_node.children:
                if param.type == 'identifier':
                    params.append({'name': self._get_node_text(param, content)})
                elif param.type == 'typed_parameter':
                    name = self._get_node_text(param.child_by_field_name('name'), content)
                    type_annotation = self._get_node_text(param.child_by_field_name('type'), content)
                    params.append({'name': name, 'type': type_annotation})
        return params
    
    def _get_jsdoc(self, node, content: str) -> str:
        """Extract JSDoc comments."""
        # This is a simplified version - real JSDoc parsing would be more complex
        return ""
    
    def _get_js_params(self, function_node) -> List[Dict[str, str]]:
        """Extract JavaScript function parameters."""
        params = []
        params_node = function_node.child_by_field_name('parameters')
        if params_node:
            for param in params_node.children:
                if param.type == 'identifier':
                    params.append({'name': self._get_node_text(param, content)})
        return params
    
    def _get_go_comment(self, node, content: str) -> str:
        """Extract Go comments."""
        return ""
    
    def _get_go_params(self, function_node) -> List[Dict[str, str]]:
        """Extract Go function parameters."""
        params = []
        params_node = function_node.child_by_field_name('parameters')
        if params_node:
            for param in params_node.children:
                if param.type == 'parameter_declaration':
                    name = self._get_node_text(param.child_by_field_name('name'), content)
                    param_type = self._get_node_text(param.child_by_field_name('type'), content)
                    params.append({'name': name, 'type': param_type})
        return params
    
    def _get_java_comment(self, node, content: str) -> str:
        """Extract Java comments."""
        return ""
    
    def _get_java_params(self, method_node) -> List[Dict[str, str]]:
        """Extract Java method parameters."""
        params = []
        params_node = method_node.child_by_field_name('parameters')
        if params_node:
            for param in params_node.children:
                if param.type == 'formal_parameter':
                    name = self._get_node_text(param.child_by_field_name('name'), content)
                    param_type = self._get_node_text(param.child_by_field_name('type'), content)
                    params.append({'name': name, 'type': param_type})
        return params